import random
import os
from glob import glob
import importlib
from src import const
import argparse
import pandas as pd
import numpy as np
import torch


def merge_const(module_name):
    new_conf = importlib.import_module(module_name)
    for key, value in new_conf.__dict__.items():
        if not(key.startswith('_')):
            # const.__dict__[key] = value
            setattr(const, key, value)
            print('override', key, value)


def parse_args_and_merge_const():
    parser = argparse.ArgumentParser()
    parser.add_argument('--conf', default='', type=str)
    args = parser.parse_args()
    if args.conf != '':
        merge_const(args.conf)


def get_train_test_df():
    test_uid = ['U02', 'U04', 'U08', 'U13', 'U17', 'U19', 'U26', 'U29']
    df = pd.read_csv('./data/info.csv')
    df['uid'] = df.Video.map(lambda x: x[:3])
    df = df[df.CroppedPath.notnull()]
    test_df = df[df.uid.isin(test_uid)]
    train_df = df[~(df.uid.isin(test_uid))]
    return train_df, test_df


def get_train_test_df_with_in_use():
    test_uid = ['U02', 'U04', 'U08', 'U13', 'U17', 'U19', 'U26', 'U29']
    df = pd.read_csv('./data/info.csv')
    df['uid'] = df.Video.map(lambda x: x[:3])
    df = df[df.CroppedPath.notnull()]
    test_df = df[df.uid.isin(test_uid)]
    train_df = pd.read_csv('./data/train_with_in_use.csv')
    return train_df, test_df


class Evaluator(object):

    def __init__(self, test_df):
        '''
        注意：务必保证test_dataloader的shuffle=False
        '''
        self.test_df = test_df
        self.reset()

    def reset(self):
        self.image_count = 0
        self.image_correct = 0
        self.all_prob = []

    def add(self, sample, output):
        _, predicted = torch.max(output.data, 1)
        self.image_count += sample['label'].size(0)
        self.image_correct += (predicted == sample['label']).sum().item()
        self.all_prob.append(output.detach().cpu().numpy())

    def rule_pred(self, all_df, fatigue_th=0.3, yawn_th=5):
        def rule_pred_func(x):
            label = str(x.Label.iloc[0])
            x = x[['ImageID', 'pred_label']].sort_values('ImageID')
            fatigue_rate = len(x[x.pred_label == 'fatigue']) / len(x)
            pred_frames = x.pred_label.tolist()
            yawn_s = 0
            max_yawn_s = 0
            for i in range(len(pred_frames)):
                if pred_frames[i] == 'yawn':
                    yawn_s += 1
                else:
                    yawn_s = 0
                if yawn_s > max_yawn_s:
                    max_yawn_s = yawn_s

            if max_yawn_s >= yawn_th:
                rule_pred_label = 'yawn'
            elif fatigue_rate >= fatigue_th:
                rule_pred_label = 'fatigue'
            else:
                rule_pred_label = 'normal'
            ret = pd.Series([
                label,
                len(x[x.pred_label == 'normal']) / len(x),
                len(x[x.pred_label == 'yawn']) / len(x),
                len(x[x.pred_label == 'fatigue']) / len(x),
                max_yawn_s,
                rule_pred_label,
            ],
                index=[
                    'label',
                    'normal_rate',
                    'yawn_rate',
                    'fatigue_rate',
                    'max_yawn_s',
                    'rule_pred_label',
            ])
            return ret
        self.rule_pred_df = all_df.groupby('VideoID').apply(rule_pred_func)
        return len(self.rule_pred_df[self.rule_pred_df['rule_pred_label'] == self.rule_pred_df['label']]) / len(self.rule_pred_df)

    def evaluate(self):
        image_accuracy = self.image_correct / self.image_count
        # softmax
        all_prob = np.concatenate(self.all_prob)
        all_prob = np.exp(all_prob)
        all_prob = all_prob / all_prob.sum(axis=1).reshape(all_prob.shape[0], -1)

        test_info = self.test_df[['VideoID', 'Label', 'CroppedPath']].copy()
        video_label = test_info.groupby('VideoID').head(1)

        test_info['pred_0'] = all_prob[:, 0]
        test_info['pred_1'] = all_prob[:, 1]
        test_info['pred_2'] = all_prob[:, 2]

        test_mean_info = test_info.groupby('VideoID').mean()
        pred_info = test_mean_info.apply(lambda x: const.label_id2name[np.argmax(x.values)],
                                         axis=1)

        pred_info = pd.DataFrame(pred_info, columns=['pred_label'])

        pred_info = video_label.set_index('VideoID').join(pred_info)

        self.pred_info = pred_info
        self.test_info = test_info

        all_info = self.test_df.copy()
        all_info['pred_0'] = all_prob[:, 0]
        all_info['pred_1'] = all_prob[:, 1]
        all_info['pred_2'] = all_prob[:, 2]
        all_info['pred_label'] = all_info[['pred_0', 'pred_1', 'pred_2']]\
            .apply(lambda x: const.label_id2name[np.argmax(x.values)], axis=1)
        self.all_info = all_info

        video_accuracy = sum(pred_info['pred_label'] == pred_info['Label']) / len(pred_info)
        ret = {
            'image_accuracy': image_accuracy,
            'video_accuracy': video_accuracy,
        }
        for i in [0.2, 0.3, 0.4]:
            for j in [3, 5, 7]:
                i_str = '%.2f' % i
                ret['rule_video_accuracy_fati{}_yawn{}'.format(i_str, j)] = \
                    self.rule_pred(all_info, i, j)
        return ret
